package cn.bugstack.openai.executor.model.chatglm.valobj;

import cn.bugstack.openai.exception.OpenAiSdkException;
import cn.bugstack.openai.util.JacksonUtil;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * ChatGLM 请求参数
 *
 * @author 小傅哥，微信：fustack
 */
@Slf4j
@Data
@JsonInclude(JsonInclude.Include.NON_NULL)
@Builder
@NoArgsConstructor
@AllArgsConstructor
public class ChatGLMCompletionRequest {

    /**
     * 模型
     */
    @Builder.Default
    private String model = Model.GLM_3_TURBO.getCode();

    /**
     * 请求ID
     */
    @JsonProperty("request_id")
    @Builder.Default
    private String requestId = String.format("xfg-%d", System.currentTimeMillis());
    /**
     * do_sample 为 true 时启用采样策略，do_sample 为 false 时采样策略 temperature、top_p 将不生效
     * 24年1月发布的 GLM_3_5_TURBO、GLM_4 模型时新增
     */
    @JsonProperty("do_sample")
    @Builder.Default
    private boolean doSample = true;
    /**
     * 控制温度【随机性】
     * 取值范围是：(0.0,1.0]，不能等于 0，默认值为 0.95,值越大，会使输出更随
     */
    @Builder.Default
    private float temperature = 0.9f;
    /**
     * 多样性控制；
     * 取值范围是：(0.0, 1.0) 开区间，不能等于 0 或 1，默认值为 0.7
     */
    @JsonProperty("top_p")
    @Builder.Default
    private float topP = 0.7f;
    /**
     * 是否流式返回
     * 使用同步调用时，此参数应当设置为 fasle 或者省略。表示模型生成完所有内容后一次性返回所有内容。
     * 如果设置为 true，模型将通过标准 Event Stream ，逐块返回模型生成内容。Event Stream 结束时会返回一条data: [DONE]消息。
     */
    @Builder.Default
    private boolean stream = true;
    /**
     * 输入给模型的会话信息
     * 用户输入的内容；role=user
     * 挟带历史的内容；role=assistant
     */
    private List<Prompt> messages;
    /**
     * 智普AI sse 固定参数 incremental = true 【增量返回】
     */
    @Builder.Default
    private boolean incremental = true;
    /**
     * sseformat, 用于兼容解决sse增量模式okhttpsse截取data:后面空格问题, [data: hello]。只在增量模式下使用sseFormat。
     */
    @Builder.Default
    private String sseFormat = "data";
    /**
     * 模型输出最大tokens
     * 24年1月发布的 GLM_3_5_TURBO、GLM_4 模型时新增
     */
    @JsonProperty("max_tokens")
    @Builder.Default
    private Integer maxTokens = 2048;
    /**
     * 模型在遇到stop所制定的字符时将停止生成，目前仅支持单个停止词，格式为["stop_word1"]
     * 24年1月发布的 GLM_3_5_TURBO、GLM_4 模型时新增
     */
    private List<String> stop;
    /**
     * 可供模型调用的工具列表,tools字段会计算 tokens ，同样受到tokens长度的限制
     * 24年1月发布的 GLM_3_5_TURBO、GLM_4 模型时新增
     */
    private List<Tool> tools;
    /**
     * 用于控制模型是如何选择要调用的函数，仅当工具类型为function时补充。默认为auto，当前仅支持auto。
     * 24年1月发布的 GLM_3_5_TURBO、GLM_4 模型时新增
     */
    @JsonProperty("tool_choice")
    @Builder.Default
    private String toolChoice = "auto";

    @Data
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class Prompt {
        private String role;
        private String content;
    }

    @Data
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class Tool {
        private Type type;
        @JsonInclude(JsonInclude.Include.NON_NULL)
        private Function function;
        private Retrieval retrieval;
        @JsonProperty("web_search")
        private WebSearch webSearch;

        public String getType() {
            return type.code;
        }

        @Getter
        @AllArgsConstructor
        public enum Type {
            function("function", "函数功能"),
            retrieval("retrieval", "知识库"),
            web_search("web_search", "联网"),
            ;
            private final String code;
            private final String info;
        }

        @Data
        @Builder
        @NoArgsConstructor
        @AllArgsConstructor
        public static class Function {
            // 函数名称，只能包含a-z，A-Z，0-9，下划线和中横线。最大长度限制为64
            private String name;
            // 用于描述函数功能。模型会根据这段描述决定函数调用方式。
            private String description;
            // parameter 字段需要传入一个 Json Schema 对象，以准确地定义函数所接受的参数。https://open.bigmodel.cn/dev/api#glm-3-turbo
            private Object parameters;
        }

        @Data
        @Builder
        @NoArgsConstructor
        @AllArgsConstructor
        public static class Retrieval {
            // 当涉及到知识库ID时，请前往开放平台的知识库模块进行创建或获取。
            @JsonProperty("knowledge_id")
            private String knowledgeId;
            // 请求模型时的知识库模板，默认模板：
            @JsonProperty("prompt_template")
            @Builder.Default
            private String promptTemplate = "\"\"\"\n" +
                    "{{ knowledge}}\n" +
                    "\"\"\"\n" +
                    "中找问题\n" +
                    "\"\"\"\n" +
                    "{{question}}\n" +
                    "\"\"\"";
        }

        // 仅当工具类型为web_search时补充，如果tools中存在类型retrieval，此时web_search不生效。
        @Data
        @Builder
        @NoArgsConstructor
        @AllArgsConstructor
        public static class WebSearch {
            // 是否启用搜索，默认启用搜索 enable = true/false
            @Builder.Default
            private Boolean enable = true;
            // 强制搜索自定义关键内容，此时模型会根据自定义搜索关键内容返回的结果作为背景知识来回答用户发起的对话。
            @JsonProperty("search_query")
            private String searchQuery;
        }


    }

    @Override
    public String toString() {
        Map<String, Object> paramsMap = new HashMap<>();
        paramsMap.put("model", this.model);
        if (null == this.messages) {
            throw new OpenAiSdkException(400, "One of messages or prompt must not be empty！");
        }
        paramsMap.put("messages", this.messages);
        if (null != this.requestId) {
            paramsMap.put("request_id", this.requestId);
        }
        paramsMap.put("do_sample", this.doSample);

        paramsMap.put("stream", this.stream);
        paramsMap.put("temperature", this.temperature);
        paramsMap.put("top_p", this.topP);
        paramsMap.put("max_tokens", this.maxTokens);
        if (null != this.stop && !this.stop.isEmpty()) {
            paramsMap.put("stop", this.stop);
        }
        if (null != this.tools && !this.tools.isEmpty()) {
            paramsMap.put("tools", this.tools);
            paramsMap.put("tool_choice", this.toolChoice);
        }
        return JacksonUtil.toJsonString(paramsMap);
    }
}
